import numpy as np
from PIL import Image
import torch
from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
import torchvision.transforms.functional as F

class TransformTwice:
    def __init__(self, transform):
        self.transform = transform

    def __call__(self, inp):
        out1 = self.transform(inp)
        out2 = self.transform(inp)
        return out1, out2

class SimpleDataset(Dataset):
    def __init__(self, data, target, transform):
        self.data = data
        self.target = target
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img_tensor = self.data[index]
        sample = self.transform(img_tensor)
        # 当为无标记样本数据集时，
        if self.target is not None:
            target = self.target[index]
        else:
            target = -1
        return sample, target

class RotationDataset(Dataset):
    def __init__(self, data, transform):
        self.data = data
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        img_tensor = self.data[index]
        transformed_tenosr = self.transform(img_tensor)
        rotated_imgs = [
            transformed_tenosr,
            F.rotate(transformed_tenosr, 90),
            F.rotate(transformed_tenosr, 180),
            F.rotate(transformed_tenosr, 270)
            ]
        rotation_labels = torch.LongTensor([0, 1, 2, 3])
        # 当为无标记样本数据集时，
        return torch.stack(rotated_imgs, dim=0), rotation_labels

class TripletDataset(Dataset):
    def __init__(self, data, transform, triplet_sampler):
        self.data = data
        self.transform = transform
        self.triplet_sampler = triplet_sampler

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        anchor_idx, positive_idx, negetive_idx = self.triplet_sampler(index)
        # pull data
        anchor_tensors = self.data[anchor_idx]
        positive_tensors = self.data[positive_idx]
        negative_tensors = self.data[negetive_idx]
        # transforms
        anchor_tensors = self.transform(anchor_tensors)
        positive_tensors = self.transform(positive_tensors)
        negative_tensors = self.transform(negative_tensors)
        return anchor_tensors, positive_tensors, negative_tensors


class TripletSampler(object):
    def __init__(self, targets):
        self.targets = targets
        self.unique_labels = torch.unique(targets)

    def __call__(self, index):
        propose_positive_mask = (self.targets == self.targets[index])
        # 去除 anchor 
        propose_positive_mask[index] = False
        propose_positive_idxs = propose_positive_mask.nonzero(as_tuple=True)[0]
        propose_negative_mask = (self.targets != self.targets[index])
        propose_negative_idxs = propose_negative_mask.nonzero(as_tuple=True)[0]
        positive_idx = torch.randint(0, len(propose_positive_idxs), (1,))
        negative_idx = torch.randint(0, len(propose_negative_idxs), (1,))
        return index, positive_idx.item(), negative_idx.item()